/*
 * This is written by Lulu Han and Yunguo Guan.
 * E-mail: locomotive_crypto@163.com
 */

#include "she.h"

SHE::PrivateKey::PrivateKey(const NTL::ZZ& p, const NTL::ZZ& q, const NTL::ZZ& L){
	m_N = p*q;
	m_p = p;
	m_q = q;
	m_L = L;
}

NTL::ZZ SHE::PrivateKey::get_N() const{
	return m_N;
}

NTL::ZZ SHE::PrivateKey::get_p() const{
	return m_p;
}

NTL::ZZ SHE::PrivateKey::get_q() const{
	return m_q;
}

NTL::ZZ SHE::PrivateKey::get_L() const{
	return m_L;
}

///////////////////////////////////////////////////////////////////////////////
SHE::PublicKey::PublicKey(const NTL::ZZ& N, const NTL::ZZ& e01, const NTL::ZZ& e02, const NTL::ZZ& e1, 
	const NTL::ZZ& en1, long k0, long k1, long k2){
	m_N = N;
	m_enc_01 = e01;
	m_enc_02 = e02;
	m_enc_1 = e1;
	m_enc_neg_1 = en1;
	
	m_k0 = k0;
	m_k1 = k1;
	m_k2 = k2;
}

NTL::ZZ SHE::PublicKey::get_N() const{
	return m_N;
}

long SHE::PublicKey::get_k2() const{
	return m_k2;
}

long SHE::PublicKey::get_k1() const{
	return m_k1;
}

long SHE::PublicKey::get_mul_depth() const{
	return floor(double(m_k0)/(2*m_k2) - 1);
}

NTL::ZZ SHE::PublicKey::get_enc01() const{
	return m_enc_01;
}

NTL::ZZ SHE::PublicKey::get_enc02() const{
	return m_enc_02;
}
NTL::ZZ SHE::PublicKey::get_enc1() const{
	return m_enc_1;
}

NTL::ZZ SHE::PublicKey::get_encn1() const{
	return m_enc_neg_1;
}
///////////////////////////////////////////////////////////////////////////////

SHE::Encryptor::Encryptor(const PublicKey& pk):m_pk(pk){
	// ...
}

void SHE::Encryptor::encrypt(const NTL::ZZ& pt, NTL::ZZ& ct){
	NTL::ZZ r1, r2;
	
	r1 = NTL::RandomLen_ZZ(m_pk.get_k2()); 
	r2 = NTL::RandomLen_ZZ(m_pk.get_k2()); 
	
	if( pt > 0 ){
		ct = (r1*m_pk.get_enc01() + r2*m_pk.get_enc02() + pt*m_pk.get_enc1()) % m_pk.get_N();
	}
	else if ( pt < 0 ){
		ct = (r1*m_pk.get_enc01() + r2*m_pk.get_enc02() + NTL::abs(pt)*m_pk.get_encn1()) % m_pk.get_N();
	}
	else{
		ct = (r1*m_pk.get_enc01() + r2*m_pk.get_enc02()) % m_pk.get_N();
	}
}
		
void SHE::Encryptor::he_add_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_k2() - 1);
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	
	ct = (c + ct1 + ct2) % m_pk.get_N();
}

void SHE::Encryptor::he_add_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_k2() - 1);
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	
	ct = (c + ct1 + pt1) % m_pk.get_N();
}
		
void SHE::Encryptor::he_mul_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_k2() - 1);
	
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	ct = (c + ct1 * pt1) % m_pk.get_N();
}

void SHE::Encryptor::he_mul_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_k2() - 1);
	
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	ct = (c + ct1 * ct2) % m_pk.get_N();
}

///////////////////////////////////////////////////////////////////////////////
SHE::Decryptor::Decryptor(const PrivateKey& sk):m_sk(sk){
	// ...
}

void SHE::Decryptor::decrypt(const NTL::ZZ& ct, NTL::ZZ& pt){
	NTL::ZZ t;
	
	t = (ct % m_sk.get_p()) % m_sk.get_L();
	
	if( t < (m_sk.get_L() / 2)){
		pt = t;
	}
	else{
		pt = t - m_sk.get_L();
	}
}

///////////////////////////////////////////////////////////////////////////////

void SHE::key_gen(PrivateKey& sk, PublicKey& pk, long k0, long k1, long k2){
	
	if ( k2 >= (k0 / 2 - 1)){ // throw an exception
		throw std::invalid_argument("invalid_argument: k0 > 2*k2");
	}
	
	// Generate random values p and q
	NTL::ZZ p, q, N;
	p =  NTL::GenPrime_ZZ(k0, 80);
	q =  NTL::GenPrime_ZZ(k0, 80);
	N= p*q;
	
	NTL::ZZ L;
	L = NTL::GenPrime_ZZ(k2, 80);
	
	// Generate the ciphertexts of messages 0, 1, and -1
	NTL::ZZ e01, e02, e1, en1;
	NTL::ZZ rl, rp;
	
	rl = NTL::RandomLen_ZZ(k2);
	rp = NTL::RandomLen_ZZ(k0);
	e01 = ((rl*L + 0) * (1 + rp*p)) % N;	

	rl = NTL::RandomLen_ZZ(k2);
	rp = NTL::RandomLen_ZZ(k0);
	e02 = ((rl*L + 0) * (1 + rp*p)) % N;

	rl = NTL::RandomLen_ZZ(k2);
	rp = NTL::RandomLen_ZZ(k0);
	e1 = ((rl*L + 1) * (1 + rp*p)) % N;
	
	rl = NTL::RandomLen_ZZ(k2);
	rp = NTL::RandomLen_ZZ(k0);
	en1 = ((rl*L + (L-1)) * (1 + rp*p)) % N;
	
	PrivateKey prikey(p, q, L);
	PublicKey pubkey(N, e01, e02, e1, en1, k0, k1, k2);
	
	sk = prikey;
	pk = pubkey;
}
